#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <limits.h>

#define MAX_CHAR 256 // 最大字符数（ASCII范围）
#define MAX_BIT 1000 // 最大编码长度

// 哈夫曼树结点结构
typedef struct HuffmanNode
{
    char character;            // 字符
    int frequency;             // 频率
    struct HuffmanNode *left;  // 左孩子
    struct HuffmanNode *right; // 右孩子
} HuffmanNode;

// 森林结构（用于构建哈夫曼树）
typedef struct
{
    HuffmanNode **trees; // 树指针数组
    int size;            // 当前森林大小
    int capacity;        // 森林容量
} Forest;

// 哈夫曼编码结构
typedef struct
{
    char character;     // 字符
    int frequency;      // 频率（添加缺失的成员）
    char code[MAX_BIT]; // 编码字符串
} HuffmanCode;

// 1. 创建新结点
HuffmanNode *createNode(char character, int frequency)
{
    HuffmanNode *node = (HuffmanNode *)malloc(sizeof(HuffmanNode));
    node->character = character;
    node->frequency = frequency;
    node->left = node->right = NULL;
    return node;
}

// 2. 初始森林
void initForest(Forest *forest, int capacity)
{
    forest->trees = (HuffmanNode **)malloc(capacity * sizeof(HuffmanNode *));
    forest->size = 0;
    forest->capacity = capacity;
}

// 3. 销毁哈夫曼树
void destroyHuffmanTree(HuffmanNode *root)
{
    if (root)
    {
        destroyHuffmanTree(root->left);
        destroyHuffmanTree(root->right);
        free(root);
    }
}

// 4. 销毁森林
void destroyForest(Forest *forest)
{
    for (int i = 0; i < forest->size; i++)
    {
        // 递归销毁树
        destroyHuffmanTree(forest->trees[i]);
    }
    free(forest->trees);
    forest->trees = NULL;
    forest->size = 0;
    forest->capacity = 0;
}

// 5. 添加树到森林
void addTreeToForest(Forest *forest, HuffmanNode *tree)
{
    if (forest->size >= forest->capacity)
    {
        // 扩容
        forest->capacity *= 2;
        forest->trees = (HuffmanNode **)realloc(forest->trees, forest->capacity * sizeof(HuffmanNode *));
    }
    forest->trees[forest->size++] = tree;
}

// 6. 找到频率最小的树
int findMinTree(Forest *forest)
{
    int minIndex = -1;
    int minFreq = INT_MAX;

    for (int i = 0; i < forest->size; i++)
    {
        if (forest->trees[i] && forest->trees[i]->frequency < minFreq)
        {
            minFreq = forest->trees[i]->frequency;
            minIndex = i;
        }
    }
    return minIndex;
}

// 7. 移除并返回最小树
HuffmanNode *removeMinTree(Forest *forest)
{
    int minIndex = findMinTree(forest);
    if (minIndex == -1)
        return NULL;

    HuffmanNode *minTree = forest->trees[minIndex];
    // 用最后一个元素替换
    forest->trees[minIndex] = forest->trees[forest->size - 1];
    forest->size--;
    return minTree;
}

// 8. 构建哈夫曼树
HuffmanNode *buildHuffmanTree(Forest *forest)
{
    // 构建哈夫曼树
    while (forest->size > 1)
    {
        // 取出两个最小的树
        HuffmanNode *left = removeMinTree(forest);
        HuffmanNode *right = removeMinTree(forest);

        // 创建新树（内部结点，字符设为'\0'）
        HuffmanNode *merged = createNode('\0', left->frequency + right->frequency);
        merged->left = left;
        merged->right = right;

        // 添加回森林
        addTreeToForest(forest, merged);
    }

    // 返回最后的树（哈夫曼树）
    return forest->trees[0];
}

// 9. 构建哈夫曼编码表
void buildCodeTable(HuffmanNode *root, HuffmanCode *codeTable, char *code, int depth)
{
    if (!root)
        return;

    // 叶子结点
    if (root->character != '\0')
    {
        // 复制编码和频率
        strncpy(codeTable[(int)root->character].code, code, depth);
        codeTable[(int)root->character].code[depth] = '\0';
        codeTable[(int)root->character].character = root->character;
        codeTable[(int)root->character].frequency = root->frequency;
        return;
    }

    // 左分支添加'0'
    if (root->left)
    {
        code[depth] = '0';
        buildCodeTable(root->left, codeTable, code, depth + 1);
    }

    // 右分支添加'1'
    if (root->right)
    {
        code[depth] = '1';
        buildCodeTable(root->right, codeTable, code, depth + 1);
    }
}

// 10. 打印哈夫曼编码表
void printCodeTable(HuffmanCode *codeTable)
{
    printf("Huffman Code Table:\n");
    printf("Character | Frequency | Code\n");
    printf("----------|-----------|------\n");

    for (int i = 0; i < MAX_CHAR; i++)
    {
        if (codeTable[i].code[0] != '\0')
        {
            if (codeTable[i].character == ' ')
            {
                printf("' '      | %-9d | %s\n", codeTable[i].frequency, codeTable[i].code);
            }
            else if (codeTable[i].character == '\n')
            {
                printf("\\n       | %-9d | %s\n", codeTable[i].frequency, codeTable[i].code);
            }
            else
            {
                printf("%-8c | %-9d | %s\n", codeTable[i].character, codeTable[i].frequency, codeTable[i].code);
            }
        }
    }
}

// 11. 编码字符串
char *encodeString(const char *str, HuffmanCode *codeTable)
{
    // 计算编码后字符串长度
    int encodedLength = 0;
    for (int i = 0; str[i] != '\0'; i++)
    {
        encodedLength += strlen(codeTable[(int)str[i]].code);
    }

    // 分配内存
    char *encodedStr = (char *)malloc(encodedLength + 1);
    encodedStr[0] = '\0';

    // 拼接编码
    for (int i = 0; str[i] != '\0'; i++)
    {
        strcat(encodedStr, codeTable[(int)str[i]].code);
    }

    return encodedStr;
}

// 12. 解码字符串
char *decodeString(HuffmanNode *root, const char *encodedStr)
{
    // 分配内存
    char *decodedStr = (char *)malloc(strlen(encodedStr) + 1);
    int index = 0;

    HuffmanNode *current = root;
    for (int i = 0; encodedStr[i] != '\0'; i++)
    {
        if (encodedStr[i] == '0')
        {
            current = current->left;
        }
        else
        {
            current = current->right;
        }

        // 到达叶子结点
        if (current->character != '\0')
        {
            decodedStr[index++] = current->character;
            current = root; // 回到根结点
        }
    }
    decodedStr[index] = '\0';

    return decodedStr;
}

// 13. 统计字符频率
void calculateFrequencies(const char *str, int *frequencies)
{
    for (int i = 0; str[i] != '\0'; i++)
    {
        frequencies[(int)str[i]]++;
    }
}

int main()
{
    const char *inputStr = "Hello World, This is a Huffman Tree";
    printf("Original String: %s\n", inputStr);

    // 1. 统计字符频率
    int frequencies[MAX_CHAR] = {0};
    calculateFrequencies(inputStr, frequencies);

    // 2. 创建森林
    Forest forest;
    initForest(&forest, MAX_CHAR);

    // 3. 为每个字符创建树并加入森林
    for (int i = 0; i < MAX_CHAR; i++)
    {
        if (frequencies[i] > 0)
        {
            HuffmanNode *node = createNode((char)i, frequencies[i]);
            addTreeToForest(&forest, node);
        }
    }

    // 4. 构建哈夫曼树
    HuffmanNode *huffmanTree = buildHuffmanTree(&forest);

    // 5. 构建编码表
    HuffmanCode codeTable[MAX_CHAR] = {{0}};
    char tempCode[MAX_BIT];
    buildCodeTable(huffmanTree, codeTable, tempCode, 0);

    // 6. 打印编码表
    printCodeTable(codeTable);

    // 7. 编码字符串
    char *encodedStr = encodeString(inputStr, codeTable);
    printf("\nEncoded String: %s\n", encodedStr);

    // 8. 解码字符串
    char *decodedStr = decodeString(huffmanTree, encodedStr);
    printf("Decoded String: %s\n", decodedStr);

    // 9. 验证
    if (strcmp(inputStr, decodedStr) == 0)
    {
        printf("Decoding successful!\n");
    }
    else
    {
        printf("Decoding failed!\n");
    }

    // 10. 清理内存
    destroyHuffmanTree(huffmanTree);
    destroyForest(&forest);
    free(encodedStr);
    free(decodedStr);

    return 0;
}